package zookeeper.lock;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;

/**
 * @author Arnold Yand
 * @summary 基于zookeeper的分布式锁
 * @time 2018/4/15 13:02
 */
public class ZookeeperDistributedLock implements Lock, Watcher {

  private static Logger logger = LoggerFactory.getLogger(ZookeeperDistributedLock.class);

  private ZooKeeper zk;
  private String rootPath = "/etl"; // default path
  private String lockPath = "/lock_";
  private int sessionTimeout = 3000; // default timeout 3s

  private String currentNode;//当前node
  private String preNode;//前一个node
  private CountDownLatch latch;//计数器
  private CountDownLatch connectedSignal = new CountDownLatch(1);

  public ZookeeperDistributedLock(String zkConfig, String rootPath, int sessionTimeout) throws Exception {
    if (rootPath == null) {
      rootPath = this.rootPath;
    }
    this.rootPath = rootPath;
    this.sessionTimeout = sessionTimeout;

    // 初始化与zookeeper server的连接 && 创建根节点
    zk = new ZooKeeper(zkConfig, this.sessionTimeout, this);
    Stat stat = zk.exists(this.rootPath, false); // 创建根节点，不注册监听器
    if (stat == null) {
      // 创建根节点
      zk.create(this.rootPath, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
    }
  }

  @Override
  public void lock() {
    try {
      if (this.tryLock()) {
        logger.info("node {} get lock.", currentNode);
        return;
      } else {
        waitForLock(preNode, sessionTimeout);//等待锁
      }
    } catch (Exception e) {
      logger.error(e.getMessage(), e);
      e.printStackTrace();
    }
  }

  private boolean waitForLock(String lower, long waitTime) throws InterruptedException, KeeperException {
    String path = rootPath + "/" + lower;
    Stat stat = zk.exists(path, true);//同时注册监听。
    //判断比自己小一个数的节点是否存在,如果不存在则无需等待锁,同时注册监听
    if (stat != null) {
      logger.info("waiting for {}.", path);
      this.latch = new CountDownLatch(1);
      this.latch.await(waitTime, TimeUnit.MILLISECONDS);//等待，这里应该一直等待其他线程释放锁
      this.latch = null;
    }
    return true;
  }

  @Override
  public void lockInterruptibly() throws InterruptedException {
    this.lock();
  }

  @Override
  public boolean tryLock() {
    //创建临时子节点
    try {
      currentNode = zk.create(rootPath + lockPath, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
      logger.info("node {} create successful.", currentNode);
      //取出所有子节点
      List<String> subNodes = zk.getChildren(rootPath, false);
      Collections.sort(subNodes);

      //如果是最小的节点,则表示取得锁
      if (currentNode.equals(rootPath + "/" + subNodes.get(0))) {
        logger.info("node {} is smallest.", currentNode);
        return true;
      }
      //如果不是最小的节点，找到比自己小1的节点
      String subcurrentNode = currentNode.substring(currentNode.lastIndexOf("/") + 1);
      preNode = subNodes.get(Collections.binarySearch(subNodes, subcurrentNode) - 1);//找到前一个子节点
    } catch (KeeperException e) {
      logger.error(e.getMessage(), e);
    } catch (InterruptedException e) {
      logger.error(e.getMessage(), e);
    }

    return false;
  }

  @Override
  public boolean tryLock(long waitTime, TimeUnit timeUnit) throws InterruptedException {
    try {
      if (this.tryLock()) {
        return true;
      }
      return waitForLock(preNode, waitTime);
    } catch (Exception e) {
      logger.error(e.getMessage(), e);
    }
    return false;
  }

  @Override
  public void unlock() {
    try {
      logger.info("node {} unlock.", currentNode);
      zk.delete(currentNode, -1);
      currentNode = null;
      zk.close();
    } catch (Exception e) {

    }
  }

  @Override
  public Condition newCondition() {
    return null;
  }

  @Override
  public void process(WatchedEvent watchedEvent) {
    //建立连接用
    if (watchedEvent.getState() == Event.KeeperState.SyncConnected) {
      connectedSignal.countDown();
      return;
    }
    //其他线程放弃锁的标志
    if (this.latch != null) {
      this.latch.countDown();
    }
  }

  public ZooKeeper getZk() {
    return zk;
  }

  public void setZk(ZooKeeper zk) {
    this.zk = zk;
  }

  public int getSessionTimeout() {
    return sessionTimeout;
  }

  public void setSessionTimeout(int sessionTimeout) {
    this.sessionTimeout = sessionTimeout;
  }
}
